# -*- coding: utf-8 -*-
# @Create On    : 2020/4/26 13:16
# @Author  : Bao Linfeng
# @Function    : models
# @File    : models.py
# @Software: PyCharm

import torch.nn as nn


class ConvNet(nn.Module):
    def __init__(self):
        super().__init__()
        # 1*1*28*28
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=6, kernel_size=5, padding=2, stride=1),  # 1 * 6 * 224 * 224
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2),  # 1*6*112*112
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=6, out_channels=16, kernel_size=5, padding=0, stride=1),  # 1* 16 * 56 * 56
            nn.ReLU(),
            nn.MaxPool2d(kernel_size=2)  # 1 * 16 * 28 * 28
        )
        self.fc = nn.Sequential(
            nn.Linear(1 * 16 * 28 * 28, 120),
            nn.ReLU(),
            nn.Linear(120, 84),
            nn.ReLU(),
            nn.Linear(84, 10),
            nn.LogSoftmax(dim=1)
        )

    def forward(self, x):
        in_size = x.size(0)
        out = self.conv1(x)
        out = self.conv2(out)
        out = out.view(in_size, -1)  # 1 * 2000
        out = self.fc(out)
        return out


class ConvNet2(nn.Module):
    def __init__(self):
        super().__init__()
        # 1*1*28*28
        self.conv1 = nn.Sequential(

            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=11, padding=2, stride=64),  # 1 * 16 * 223 * 223
            nn.BatchNorm2d(16),
            # nn.Dropout2d(0.2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(3)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, padding=0, stride=1),  # 3* 32 * 222 * 222
            nn.BatchNorm2d(32),
            nn.Dropout2d(0.5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)  # 1 * 32 * 111 * 111
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=0, stride=1),  # 1* 64* 110 * 110
            nn.BatchNorm2d(64),
            nn.Dropout2d(0.5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)  # 1 * 64 * 55 * 55
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=0, stride=1),  # 1* 128 * 54 * 54
            nn.BatchNorm2d(128),
            nn.Dropout2d(0.2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2, stride=2)  # 1 * 128 * 27 * 27
        )
        self.fc = nn.Sequential(
            nn.Linear(1 * 128 * 27 * 27, 1024),
            nn.ReLU(inplace=True),
            nn.Linear(1024, 128),
            nn.ReLU(inplace=True),
            nn.Linear(128, 10),
            nn.LogSoftmax(dim=1)
        )

    def forward(self, x):
        in_size = x.size(0)
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        out = out.view(in_size, -1)  # 1 * 2000
        out = self.fc(out)
        return out


class AlexNet(nn.Module):
    def __init__(self):
        super(AlexNet, self).__init__()
        self.conv1 = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=64, kernel_size=11, padding=2, stride=4),  # 1 * 16 * 223 * 223
            nn.BatchNorm2d(64),
            # nn.Dropout2d(0.2),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2)
        )
        self.conv2 = nn.Sequential(
            nn.Conv2d(in_channels=64, out_channels=192, kernel_size=5, padding=2, stride=1),  # 3* 32 * 222 * 222
            nn.BatchNorm2d(192),
            # nn.Dropout2d(0.5),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2)  # 1 * 32 * 111 * 111
        )
        self.conv3 = nn.Sequential(
            nn.Conv2d(in_channels=192, out_channels=384, kernel_size=3, padding=1, stride=1),  # 1* 64* 110 * 110
            nn.BatchNorm2d(384),
            # nn.Dropout2d(0.5),
            nn.ReLU(inplace=True),
            # nn.MaxPool2d(kernel_size=2, stride=2)  # 1 * 64 * 55 * 55
        )
        self.conv4 = nn.Sequential(
            nn.Conv2d(in_channels=384, out_channels=256, kernel_size=3, padding=1, stride=1),  # 1* 128 * 54 * 54
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2)  # 1 * 128 * 27 * 27
        )
        self.conv5 = nn.Sequential(
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=3, padding=1, stride=1),  # 1* 256 * 18 * 18
            nn.BatchNorm2d(256),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2)  # 256 * 6 * 6
        )
        self.fc = nn.Sequential(
            nn.Dropout(),
            nn.Linear(256 * 6 * 6, 4096),
            nn.ReLU(inplace=True),
            nn.Dropout(),
            nn.Linear(4096, 2048),
            nn.ReLU(inplace=True),
            nn.Linear(2048, 2),
        )

    def forward(self, x):
        in_size = x.size(0)
        out = self.conv1(x)
        out = self.conv2(out)
        out = self.conv3(out)
        out = self.conv4(out)
        out = out.view(in_size, -1)  # 1 * 2000
        out = self.fc(out)
        return out
